package trie

//6.使用 Trie 树来优化存储和查找

// TrieNode 表示Trie树的节点
type TrieNode struct {
	children map[rune]*TrieNode
	isEnd    bool
}

// Trie 表示Trie树
type Trie struct {
	root *TrieNode
}

// NewTrieNode 初始化一个新的Trie节点
func NewTrieNode() *TrieNode {
	return &TrieNode{
		children: make(map[rune]*TrieNode),
		isEnd:    false,
	}
}

// NewTrie 初始化一个新的Trie树
func NewTrie() *Trie {
	return &Trie{
		root: NewTrieNode(),
	}
}

// Insert 插入单词到Trie树中
func (t *Trie) Insert(word string) {
	node := t.root
	for _, ch := range word {
		if _, ok := node.children[ch]; !ok {
			node.children[ch] = NewTrieNode()
		}
		node = node.children[ch]
	}
	node.isEnd = true
}

// Search 查找单词是否存在于Trie树中
func (t *Trie) Search(word string) bool {
	node := t.root
	for _, ch := range word {
		if _, ok := node.children[ch]; !ok {
			return false
		}
		node = node.children[ch]
	}
	return node.isEnd
}

// StartsWith 查找是否存在以给定前缀开头的单词
func (t *Trie) StartsWith(prefix string) bool {
	node := t.root
	for _, ch := range prefix {
		if _, ok := node.children[ch]; !ok {
			return false
		}
		node = node.children[ch]
	}
	return true
}
